#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
#       createiso.py
#  
#       Author: weiketing <weikting@gmail.com>
#       
from __future__ import print_function

import subprocess
import os
import config
import sys



def call_cmd(*args,**kwargs):
    print_cmd=True
    try:
        print_cmd=kwargs['print_cmd']
    except:pass
    if print_cmd:
        if isinstance(args[0],list):
#            print('"'+'" "'.join(args[0])+'"',file=sys.stderr)
            print(' '.join(args[0]),file=sys.stderr)
        else:
            print(args[0],file=sys.stderr)
    return subprocess.check_call(*args,**kwargs)


class mkisofs(object):
    
    _mkisofs = 'mkisofs'
    _options = ['-J','-r',
                "-hide-rr-moved", "-hide-joliet-trans-tbl",
                '-input-charset','utf-8',
                '-l',
                '-udf',
                '-allow-limited-size',
                '-b', 'isolinux/isolinux.bin',
                '-c', 'isolinux/boot.cat',
                '-no-emul-boot',
                '-boot-info-table', '-boot-load-size', '4',
                ]

    def __init__(self,iso_root,label=config.CD_LABEL):
        self._iso_root = iso_root
        self._label = label
        self._efi_support = False
        try:
            self._efi_support = config.EFI_SUPPORT
        except:pass

    def _iso_exists(self,path):
        fn = os.path.join(self._iso_root,path.strip('/'))
        return os.path.exists(fn)

    @property
    def options(self):
        opts = []
        if self._efi_support and self._iso_exists('boot/macboot.img'):
            opts += ['-eltorito-alt-boot','-e','boot/efiboot.img','-no-emul-boot',
                     '-eltorito-alt-boot','-e','boot/macboot.img','-no-emul-boot'
                     ]

        if self._label and len(self._label) > 0:
            opts += ['-V',self._label]

        if os.path.exists('exclude.list'):
            opts += ['-exclude-list', 'exclude.list']

        return self._options + opts

    def _isohybrid(self,isofile):
        cmd=['isohybrid']
        if self._efi_support and self._iso_exists('boot/efiboot.img'):
            cmd += ['-u',
                    '-m'
                    ]
        cmd += [isofile]
        return call_cmd(cmd)

    def mkiso(self,isofile=None):
        if isofile is None:
            if config.ISO_FILENAME and len(config.ISO_FILENAME ) > 0:
                isofile = config.ISO_FILENAME
            else:
                isofile = self._label.lower()+'.iso'
            if not isofile.lower().endswith('.iso'):
                isofile+='.iso'
            if self._iso_root == '.' or self._iso_root == './':
                isofile = '../'+isofile
        cmd = [self._mkisofs] + self.options + ['-o',isofile,self._iso_root]
        for d in config.EXTRA_ROOTS.split():
            if os.path.isdir(d):
                cmd.append(d)
        print("Generating iso ..." ,file=sys.stderr)
        rs = call_cmd(cmd)
        self._isohybrid(isofile)
        return rs


class Grub2Common(object):

    _modules = ['echo','search','normal','part_msdos','part_gpt', 'fat', 'ntfs', 'ext2','ntfs', 'udf', 'iso9660']
    _extra_modules =[]
    
    _cfg = 'grub_init.cfg'
    _prefix = config.GRUB2_PREFIX
    _img = 'core.img'
    _format = 'i386-pc'
    _root = config.ISO_ROOT


    @property
    def cfg(self):
        return self._cfg

    def __cmdline(self):
        opts=['grub-mkimage',
            '-p',self._prefix]
        if os.path.exists(self._cfg):
            opts +=['--config',self.cfg]

        opts = opts + \
            ['-o',self._img,
              '-O',self._format] 
        
        cmd= opts + self._modules + self._extra_modules

        return cmd

    def set_root(self,_root):
        self._root = _root

    def makeimg(self):
        self.gencfg()
        cmd = self.__cmdline()
        call_cmd(cmd)
        return self._img

    def isolinux_cfg(self):
        return ""

    def gencfg(self):
        cmd='sed "s|@PREFIX@|' + self._prefix +'|g" grub2_load.cfg.in >' + self.cfg
        call_cmd(cmd,shell=True)
        if self._format == 'x86_64-efi':
            cmd='echo "set efi_mode=true\nexport efi_mode" >>' + self.cfg
            call_cmd(cmd,shell=True)

class Grub2I386(Grub2Common):
    _extra_modules = ['biosdisk']
    _cfg = 'grub2_pc_init.cfg'

    def isolinux_cfg(self):
        icfg ="label Grub2\n"
        icfg+="     menu label ^Grub2\n"
        icfg+="kernel "+os.path.join(self._prefix,self._img)
        return icfg

class Grub2Uefi(Grub2Common):

    _cfg = 'grub2_efi_init.cfg'
    _format = 'x86_64-efi'
    _img = 'grubx64.efi'
    _label = "GrubCD"

    def __init__(self):
        super(Grub2Uefi,self).__init__()

    @property
    def efiboot_img(self):
        return os.path.join(self._root,'boot/efiboot.img')

    @property
    def macboot_img(self):
        return os.path.join(self._root,'boot/macboot.img')

    def setLabel(self,_label):
        self._label = _label

    def getLabel(self):
        return self._label
    
    label = property(getLabel,setLabel)

    def _make_macboot(self):
        call_cmd('./build_macboot.sh ' + self.macboot_img + " " + self.label,shell=True)
        return self.macboot_img

    def _make_efiboot(self):
        call_cmd("mformat -C -f 2880 -L 16 -i "+self.efiboot_img+" ::",shell=True)
        call_cmd("mcopy -s -i "+self.efiboot_img+ ' '+self._root+'/EFI ::/',shell=True)
        return [self.efiboot_img,self._make_macboot()]

    def make_bootimg(self):
        return self._make_efiboot()


def strip_grub_exe_ignore_cd(grub_exe_file):
    r=None
    try:
        r=open(grub_exe_file,"r")
        grub_exe_data=r.read()
        str_head="pxe detect\nconfigfile\ndefault "
        i=grub_exe_data.find(str_head)
    except:raise
    finally:
        if r:r.close()
    #print grub_exe_data[i:]
    s=grub_exe_data[i:].replace('--ignore-cd','           ')

    w = None
    yes = False
    fn = None
    try:
        w0,fn=tempfile.mkstemp(prefix=grub_exe_file)
        os.close(w0)
        w=open(fn,"w")
        w.write(grub_exe_data[:i])
        w.write(s)
        yes = True
    except:raise
    finally:
        if w:w.close()
        if yes and fn:
            os.rename(fn,grub_exe_file)
        else:
            os.remove(grub_exe_file)

def mkiso(label=config.CD_LABEL,iso_root=config.ISO_ROOT):
    mi=mkisofs(iso_root,label)
    mi.mkiso()

def grub2boot():
    gi = Grub2I386()
    ge = Grub2Uefi()
    gi.makeimg()
    ge.makeimg()
    print(gi.isolinux_cfg())

def copy_bootloader_files(iso_root=config.ISO_ROOT):
    cmd=['bash','copybooloaderfiles']
    call_cmd(cmd)


funcs = {
    'strip_grub_ignore_cd':strip_grub_exe_ignore_cd,
    'mkiso':mkiso,
    'grub2boot':grub2boot
}

def test():
    gi = Grub2I386()
    ge = Grub2Uefi()
    gi.makeimg()
    ge.makeimg()
    print(gi.isolinux_cfg())

def main():
    import sys
    func=sys.argv[1]
    args=[]
    kwargs={}
    for o in sys.argv[2:]:
        if '=' in o:
            k,v = o.split('=',2)
            kwargs[k]=v
        else:
            args.append(o)
    callback=None
    try:
        callback = funcs[func]
    except:pass
    if callback:
        callback(*args,**kwargs)

def createiso():
    _efi_support = False
    try:
        _efi_support = config.EFI_SUPPORT
    except:pass

    gi = Grub2I386()
    ge = Grub2Uefi()
    if os.path.exists(config.ISO_ROOT):
        call_cmd('rm -r ' +config.ISO_ROOT,shell=True)
    gi.makeimg()
    ge.makeimg()
    copy_bootloader_files()
    ge._make_efiboot()
    call_cmd("bash gengrubcfg",shell=True)
    mkiso()

if __name__ == '__main__':
    createiso()
